#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
数据库初始化脚本
创建示例数据：用户、班级、题目、作业等
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from app import create_app
from models import db
from models import User, Class, Problem, Assignment, AssignmentSubmission
from datetime import datetime, timedelta
import json

def init_database():
    """初始化数据库"""
    app = create_app()
    with app.app_context():
        # 创建所有表
        db.create_all()
        print("✓ 数据库表创建完成")
        
        # 清空现有数据（按照外键依赖顺序）
        from sqlalchemy import text
        
        # 禁用外键检查
        db.session.execute(text("SET FOREIGN_KEY_CHECKS = 0"))
        
        # 清空所有表
        tables = [
            'assignment_submissions',
            'assignments', 
            'class_student',
            'classes',
            'submissions',
            'problems',
            'system_logs',
            'users'
        ]
        
        for table in tables:
            try:
                db.session.execute(text(f"TRUNCATE TABLE {table}"))
                print(f"✓ 清空表 {table}")
            except Exception as e:
                print(f"⚠ 清空表 {table} 失败: {e}")
        
        # 重新启用外键检查
        db.session.execute(text("SET FOREIGN_KEY_CHECKS = 1"))
        db.session.commit()
        print("✓ 清空现有数据完成")
        
        # 创建示例用户
        create_sample_users()
        
        # 创建示例题目
        create_sample_problems()
        
        # 创建示例班级
        create_sample_classes()
        
        # 创建示例作业
        create_sample_assignments()
        
        print("✓ 数据库初始化完成！")

def create_sample_users():
    """创建示例用户"""
    users = [
        # 管理员
        {
            'username': 'admin',
            'password': 'admin123',
            'name': '系统管理员',
            'email': 'admin@oj.com',
            'role': 'admin'
        },
        # 教师
        {
            'username': 'teacher1',
            'password': 'teacher123',
            'name': '张老师',
            'email': 'teacher1@oj.com',
            'role': 'teacher'
        },
        {
            'username': 'teacher2',
            'password': 'teacher123',
            'name': '李老师',
            'email': 'teacher2@oj.com',
            'role': 'teacher'
        },
        # 学生
        {
            'username': 'student1',
            'password': 'student123',
            'name': '张三',
            'email': 'student1@oj.com',
            'role': 'student'
        },
        {
            'username': 'student2',
            'password': 'student123',
            'name': '李四',
            'email': 'student2@oj.com',
            'role': 'student'
        },
        {
            'username': 'student3',
            'password': 'student123',
            'name': '王五',
            'email': 'student3@oj.com',
            'role': 'student'
        },
        {
            'username': 'student4',
            'password': 'student123',
            'name': '赵六',
            'email': 'student4@oj.com',
            'role': 'student'
        },
        {
            'username': 'student5',
            'password': 'student123',
            'name': '钱七',
            'email': 'student5@oj.com',
            'role': 'student'
        }
    ]
    
    for user_data in users:
        user = User(
            username=user_data['username'],
            password=user_data['password'],
            name=user_data['name'],
            email=user_data['email'],
            role=user_data['role']
        )
        db.session.add(user)
    
    db.session.commit()
    print(f"✓ 创建了 {len(users)} 个示例用户")

def create_sample_problems():
    """创建示例题目"""
    problems = [
        {
            'title': '两数之和',
            'description': '给定一个整数数组 nums 和一个整数目标值 target，请你在该数组中找出和为目标值 target 的那两个整数，并返回它们的数组下标。',
            'difficulty': 'easy',
            'time_limit': 1000,
            'memory_limit': 128,
            'input_format': '第一行输入一个整数 n，表示数组长度\n第二行输入 n 个整数，表示数组 nums\n第三行输入一个整数 target',
            'output_format': '输出两个整数，表示满足条件的两个数的下标',
            'sample_input': '4\n2 7 11 15\n9',
            'sample_output': '0 1',
            'hint': '可以使用哈希表来优化时间复杂度',
            'solution': '```python\ndef twoSum(nums, target):\n    hash_map = {}\n    for i, num in enumerate(nums):\n        complement = target - num\n        if complement in hash_map:\n            return [hash_map[complement], i]\n        hash_map[num] = i\n    return []\n```',
            'test_cases': [
                {'input': '4\n2 7 11 15\n9', 'output': '0 1'},
                {'input': '3\n3 2 4\n6', 'output': '1 2'},
                {'input': '2\n3 3\n6', 'output': '0 1'}
            ]
        },
        {
            'title': '反转字符串',
            'description': '编写一个函数，其作用是将输入的字符串反转过来。输入字符串以字符数组 char[] 的形式给出。',
            'difficulty': 'easy',
            'time_limit': 1000,
            'memory_limit': 128,
            'input_format': '输入一行字符串',
            'output_format': '输出反转后的字符串',
            'sample_input': 'hello',
            'sample_output': 'olleh',
            'hint': '可以使用双指针法',
            'solution': '```python\ndef reverseString(s):\n    left, right = 0, len(s) - 1\n    while left < right:\n        s[left], s[right] = s[right], s[left]\n        left += 1\n        right -= 1\n    return s\n```',
            'test_cases': [
                {'input': 'hello', 'output': 'olleh'},
                {'input': 'world', 'output': 'dlrow'},
                {'input': 'python', 'output': 'nohtyp'}
            ]
        },
        {
            'title': '最大子数组和',
            'description': '给你一个整数数组 nums，请你找出一个具有最大和的连续子数组（子数组最少包含一个元素），返回其最大和。',
            'difficulty': 'medium',
            'time_limit': 1000,
            'memory_limit': 128,
            'input_format': '第一行输入一个整数 n，表示数组长度\n第二行输入 n 个整数，表示数组 nums',
            'output_format': '输出一个整数，表示最大子数组和',
            'sample_input': '5\n-2 1 -3 4 -1',
            'sample_output': '4',
            'hint': '可以使用动态规划或贪心算法',
            'solution': '```python\ndef maxSubArray(nums):\n    max_sum = current_sum = nums[0]\n    for num in nums[1:]:\n        current_sum = max(num, current_sum + num)\n        max_sum = max(max_sum, current_sum)\n    return max_sum\n```',
            'test_cases': [
                {'input': '5\n-2 1 -3 4 -1', 'output': '4'},
                {'input': '1\n1', 'output': '1'},
                {'input': '5\n5 4 -1 7 8', 'output': '23'}
            ]
        },
        {
            'title': '合并两个有序链表',
            'description': '将两个升序链表合并为一个新的升序链表并返回。新链表是通过拼接给定的两个链表的所有节点组成的。',
            'difficulty': 'easy',
            'time_limit': 1000,
            'memory_limit': 128,
            'input_format': '第一行输入链表1的长度和链表2的长度\n第二行输入链表1的节点值\n第三行输入链表2的节点值',
            'output_format': '输出合并后的链表节点值',
            'sample_input': '3 3\n1 2 4\n1 3 4',
            'sample_output': '1 1 2 3 4 4',
            'hint': '可以使用递归或迭代的方法',
            'solution': '```python\ndef mergeTwoLists(l1, l2):\n    if not l1:\n        return l2\n    if not l2:\n        return l1\n    \n    if l1.val <= l2.val:\n        l1.next = mergeTwoLists(l1.next, l2)\n        return l1\n    else:\n        l2.next = mergeTwoLists(l1, l2.next)\n        return l2\n```',
            'test_cases': [
                {'input': '3 3\n1 2 4\n1 3 4', 'output': '1 1 2 3 4 4'},
                {'input': '0 0\n\n', 'output': ''},
                {'input': '0 1\n\n1', 'output': '1'}
            ]
        },
        {
            'title': '有效的括号',
            'description': '给定一个只包括 \'(\'，\')\'，\'{\'，\'}\'，\'[\'，\']\' 的字符串 s，判断字符串是否有效。',
            'difficulty': 'easy',
            'time_limit': 1000,
            'memory_limit': 128,
            'input_format': '输入一行字符串，包含括号',
            'output_format': '输出 true 或 false',
            'sample_input': '()',
            'sample_output': 'true',
            'hint': '可以使用栈来匹配括号',
            'solution': '```python\ndef isValid(s):\n    stack = []\n    brackets = {\')\': \'(\', \'}\': \'{\', \']\': \'[\'}\n    \n    for char in s:\n        if char in \'({[\':\n            stack.append(char)\n        else:\n            if not stack or stack.pop() != brackets[char]:\n                return False\n    \n    return len(stack) == 0\n```',
            'test_cases': [
                {'input': '()', 'output': 'true'},
                {'input': '()[]{}', 'output': 'true'},
                {'input': '(]', 'output': 'false'},
                {'input': '([)]', 'output': 'false'}
            ]
        },
        {
            'title': '爬楼梯',
            'description': '假设你正在爬楼梯。需要 n 阶你才能到达楼顶。每次你可以爬 1 或 2 个台阶。你有多少种不同的方法可以爬到楼顶呢？',
            'difficulty': 'easy',
            'time_limit': 1000,
            'memory_limit': 128,
            'input_format': '输入一个整数 n，表示楼梯阶数',
            'output_format': '输出一个整数，表示不同的爬法数量',
            'sample_input': '3',
            'sample_output': '3',
            'hint': '可以使用动态规划，类似于斐波那契数列',
            'solution': '```python\ndef climbStairs(n):\n    if n <= 2:\n        return n\n    \n    dp = [0] * (n + 1)\n    dp[1] = 1\n    dp[2] = 2\n    \n    for i in range(3, n + 1):\n        dp[i] = dp[i-1] + dp[i-2]\n    \n    return dp[n]\n```',
            'test_cases': [
                {'input': '3', 'output': '3'},
                {'input': '2', 'output': '2'},
                {'input': '4', 'output': '5'}
            ]
        }
    ]
    
    for problem_data in problems:
        problem = Problem(
            title=problem_data['title'],
            description=problem_data['description'],
            type='code',  # 编程题
            difficulty=problem_data['difficulty'],
            time_limit=problem_data['time_limit'],
            memory_limit=problem_data['memory_limit'],
            language='python',
            created_by=1,  # 管理员创建
            test_cases=json.dumps(problem_data['test_cases']),
            answer=problem_data['solution']  # 使用solution作为答案
        )
        db.session.add(problem)
    
    db.session.commit()
    print(f"✓ 创建了 {len(problems)} 个示例题目")

def create_sample_classes():
    """创建示例班级"""
    # 获取教师用户
    teacher1 = User.query.filter_by(username='teacher1').first()
    teacher2 = User.query.filter_by(username='teacher2').first()
    
    # 获取学生用户
    students = User.query.filter_by(role='student').all()
    
    classes = [
        {
            'name': 'Python编程基础',
            'description': '学习Python编程的基础知识，包括变量、函数、循环、条件语句等',
            'teacher': teacher1,
            'max_students': 30,
            'students': students[:3]  # 前3个学生
        },
        {
            'name': '数据结构与算法',
            'description': '学习基本的数据结构和算法，包括数组、链表、栈、队列、排序等',
            'teacher': teacher1,
            'max_students': 25,
            'students': students[1:4]  # 中间3个学生
        },
        {
            'name': 'Web开发入门',
            'description': '学习Web开发的基础知识，包括HTML、CSS、JavaScript等',
            'teacher': teacher2,
            'max_students': 35,
            'students': students[2:]  # 后3个学生
        }
    ]
    
    for class_data in classes:
        class_obj = Class(
            name=class_data['name'],
            description=class_data['description'],
            teacher_id=class_data['teacher'].id,
            max_students=class_data['max_students'],
            status='active'
        )
        db.session.add(class_obj)
        db.session.flush()  # 获取ID
        
        # 添加学生到班级
        for student in class_data['students']:
            class_obj.students.append(student)
    
    db.session.commit()
    print(f"✓ 创建了 {len(classes)} 个示例班级")

def create_sample_assignments():
    """创建示例作业"""
    # 获取班级
    class1 = Class.query.filter_by(name='Python编程基础').first()
    class2 = Class.query.filter_by(name='数据结构与算法').first()
    class3 = Class.query.filter_by(name='Web开发入门').first()
    
    # 获取题目
    problems = Problem.query.all()
    
    assignments = [
        {
            'title': 'Python基础练习',
            'description': '完成Python基础语法练习，包括变量、函数、循环等',
            'class_obj': class1,
            'problems': problems[:2],  # 前2个题目
            'due_date': datetime.now() + timedelta(days=7),
            'total_score': 100
        },
        {
            'title': '算法基础作业',
            'description': '完成基础算法练习，包括数组操作、字符串处理等',
            'class_obj': class2,
            'problems': problems[1:4],  # 中间3个题目
            'due_date': datetime.now() + timedelta(days=5),
            'total_score': 150
        },
        {
            'title': 'Web开发实践',
            'description': '完成Web开发基础练习，包括HTML、CSS等',
            'class_obj': class3,
            'problems': problems[2:4],  # 中间2个题目
            'due_date': datetime.now() + timedelta(days=10),
            'total_score': 120
        }
    ]
    
    for assignment_data in assignments:
        problem_ids = [p.id for p in assignment_data['problems']]
        assignment = Assignment(
            title=assignment_data['title'],
            description=assignment_data['description'],
            class_id=assignment_data['class_obj'].id,
            problem_ids=json.dumps(problem_ids),
            due_date=assignment_data['due_date'],
            total_score=assignment_data['total_score'],
            status='active'
        )
        db.session.add(assignment)
    
    db.session.commit()
    print(f"✓ 创建了 {len(assignments)} 个示例作业")

if __name__ == '__main__':
    init_database()
